package com.taint.agent;


import com.taint.callgraph.TraceItem;
import com.taint.configuration.Configuration;
import com.taint.filewatch.FileWatcher;
import com.taint.sourcesink.MethodHook;
import com.taint.utils.HookType;
import com.taint.utils.Utils;
import net.bytebuddy.agent.builder.AgentBuilder;

import java.io.File;
import java.io.IOException;
import java.lang.instrument.Instrumentation;
import java.util.*;

import static com.taint.utils.HookType.FILEWATCHPATH;


/*
    Author: taochao
     */

public class FlowTaintAgent {

    private static ArrayList sinks;
    private static ArrayList sources;
    public static boolean isFirstLoaded = true;


    public static LinkedHashMap<String, Set<MethodHook>> methodSinkHookMap = new LinkedHashMap<>();
    public static LinkedHashMap<String, Set<MethodHook>> methodSourceHookMap = new LinkedHashMap<>();

    public static LinkedHashMap<String, Set<MethodHook>> getMethodSinkHookMap() {
        return methodSinkHookMap;
    }

    public static HashMap<String,Boolean> MatchedLoadedClazz=new HashMap<>();


    public static ThreadLocal<Stack<TraceItem>> stacks = new ThreadLocal() {
        @Override
        protected Stack<TraceItem> initialValue() {
            return new Stack<>();
        }
    };

    public static ThreadLocal<LinkedList<TraceItem>> calls = new ThreadLocal() {
        @Override
        protected LinkedList<TraceItem> initialValue() {
            return new LinkedList<>();
        }
    };

    public static void isFirstLoad() {
        sinks = (ArrayList) Configuration.LoadYaml().get("sinks");
        sources = (ArrayList) Configuration.LoadYaml().get("sources");
        methodSinkHookMap = MethodHook.InsertClazz(sinks, methodSinkHookMap);
        methodSourceHookMap = MethodHook.InsertClazz(sources, methodSourceHookMap);
        Utils.isExist(FILEWATCHPATH + File.separator + "methods.yaml");
    }


    public FlowTaintAgent() {

    }

    public static void agentmain(String arg, Instrumentation instrumentation) {
        //logger.info("AgentMain is loading");

        init(arg, instrumentation);

        FileWatcher watcher = new FileWatcher(new String[]{FILEWATCHPATH});
        try {
            watcher.start(
                    path -> {
                        //参数为变化的文件
                        String fileName = path.getFileName().toString();
                        if (fileName.equals("methods.yaml")) {
                            Configuration.LoadYaml(FILEWATCHPATH + File.separator + fileName);
                            //这里需要reset加载的类
                            FlowTaintClassTransformer.getTransformer().reset(instrumentation, AgentBuilder.RedefinitionStrategy.RETRANSFORMATION);
                            init(arg, instrumentation);
                        }
                    });
        } catch (IOException e) {
            e.printStackTrace();
        }

    }

    public static void premain(String arg, Instrumentation instrumentation) {
        init(arg, instrumentation);
    }

    private static void init(String arg, Instrumentation instrumentation){
        // logger.info("enter ByteBuddy!");
        if (isFirstLoaded) {
            isFirstLoad();
            isFirstLoaded = false;
        }

        try {
            FlowTaintClassLoader.injectBootstrapClasses(instrumentation);
        } catch (IOException | ClassNotFoundException e) {
            e.printStackTrace();
        }
        // logger.info("end ByteBuddy!");
        FlowTaintClassTransformer.ForwardStaticAnaylzer(instrumentation,MatchedLoadedClazz);
        FlowTaintClassTransformer.InstrumenterSourceSink(HookType.SOURCE, methodSourceHookMap, instrumentation);
        FlowTaintClassTransformer.InstrumenterSourceSink(HookType.SINK, methodSinkHookMap, instrumentation);

        //FlowTaintClassTransformer.retransform(instrumentation);
        System.out.println("size: "+MatchedLoadedClazz.size());

    }


    public static void Load() {
        Properties props = new Properties();
        try {
            props.load(FlowTaintAgent.class.getResourceAsStream("/log4j.properties"));
        } catch (IOException e) {
            e.printStackTrace();
        }
        //PropertyConfigurator.configure(props);
    }


}
